import torch
import torch.nn as nn
import torch.nn.functional as F
from . import preprocess_utils as putils
from .preprocess_utils import *

class Preprocess_Line2Window(nn.Module):
    '''
    the preprocess class for grid-with-line pipeline
    '''
    def __init__(self, configs, device=None, vis=False):
        super(Preprocess_Line2Window, self).__init__()
        self.__lossname__ = 'Preprocess_Line2Window'
        self.config = configs
        self.kps_generator = getattr(putils, self.config['kps_generator'])
        self.t_base = self.config['temperature_base']
        self.t_max = self.config['temperature_max']
        if device is not None:
            self.device = device

    def name(self):
        return self.__lossname__

    def forward(self, inputs, outputs):
        preds1 = outputs['preds1']
        preds2 = outputs['preds2']

        xf1 = preds1['local_map']
        xf2 = preds2['local_map']
        h1i, w1i = inputs['im1'].size()[2:]
        h2i, w2i = inputs['im2'].size()[2:]
        b, _, hf, wf = xf1.shape

        coord1 = inputs['coord1']
        coord2 = inputs['coord2']
        coord1_n = normalize_coords(coord1, h1i, w1i)
        coord2_n = normalize_coords(coord2, h2i, w2i)
        feat1_fine = sample_feat_by_coord(xf1, coord1_n.reshape(b, -1, 2), self.config['loss_distance'] == 'cos')
        feat2_fine = sample_feat_by_coord(xf2, coord2_n.reshape(b, -1, 2), self.config['loss_distance'] == 'cos')


        cos_sim = feat1_fine @ feat2_fine.transpose(1,2) # bxmxn
        scores0 = F.log_softmax(cos_sim, 2)
        scores1 = F.log_softmax(cos_sim.transpose(-1, -2).contiguous(), 2).transpose(-1, -2)
        scores = scores0 + scores1

        return {'scores':scores}


class Preprocess_Skip(nn.Module):
    '''
    the preprocess class for keypoint detection net training
    '''
    def __init__(self, **kargs):
        super(Preprocess_Skip, self).__init__()
        self.__lossname__ = 'Preprocess_Skip'

    def forward(self, inputs, outputs):
        return None
